import pygame
import sys
import time
import numpy as np

# ==================== 配置参数 ====================
# 基础PID参数
BASE_P = 0.8  # 基础比例系数
BASE_I = 0.15  # 基础积分系数
BASE_D = 0.3  # 基础微分系数

# 自适应参数设置
ADAPTIVE_PARAMS = {
    "large_error_threshold": 120,  # 大误差阈值（像素）
    "medium_error_threshold": 60,  # 中等误差阈值（像素）
    "large_error_factors": (1.2, 0.3, 1.5),  # (P比例, I比例, D比例)
    "medium_error_factors": (1.0, 0.8, 1.2),
    "small_error_factors": (0.6, 1.2, 1.0)
}

# 轨迹预测设置
PREDICTION_TIME = 0  # 预测时间（秒）
MAX_PREDICTION_TIME = 0.3  # 最大预测时间
MIN_PREDICTION_TIME = 0.05  # 最小预测时间

# 输出限制和抗饱和
OUTPUT_LIMIT = 15  # 最大输出速度（像素/帧）
ANTI_WINDUP_THRESHOLD = 0.85  # 抗饱和激活阈值（85%）

# 积分限幅
I_MAX = 4.0
I_MIN = -4.0

# 平滑因子
SMOOTH_FACTOR = 0.2  # 用于小球速度平滑

# 鼠标历史记录长度
MOUSE_HISTORY_SIZE = 5

# ==================== 初始化 ====================
pygame.init()
screen_size = (800, 600)
screen = pygame.display.set_mode(screen_size)
pygame.display.set_caption("高级PID控制器 - 自适应参数、轨迹预测与抗饱和")

# 字体
font = pygame.font.SysFont('Arial', 20)
large_font = pygame.font.SysFont('Arial', 24)

# 颜色定义
COLOR_BG = (240, 240, 240)  # 背景色
COLOR_BALL = (220, 20, 60)  # 红色小球
COLOR_TARGET = (30, 144, 255)  # 蓝色目标
COLOR_PREDICTION = (50, 205, 50)  # 绿色预测点
COLOR_LINE_CURRENT = (200, 200, 200)  # 当前连线
COLOR_LINE_PREDICT = (150, 200, 150)  # 预测连线
COLOR_WARNING = (255, 69, 0)  # 警告色
COLOR_SUCCESS = (34, 139, 34)  # 成功色

# 小球属性
ball_pos = [400, 300]  # 初始位置
ball_radius = 6
ball_velocity = [0, 0]  # 小球速度

# 目标标记
target_radius = 6
prediction_radius = 4

# ==================== 状态变量 ====================
# PID状态
I_past_x = 0
I_past_y = 0
previous_error_x = 0
previous_error_y = 0

# 鼠标状态
prev_mouse_pos = None
mouse_velocity = [0, 0]
mouse_acceleration = [0, 0]
mouse_history = []  # 存储最近的鼠标位置

# 性能统计
frame_count = 0
total_frames = 0
start_time = time.time()
fps_history = []
last_fps_update = time.time()

# 控制模式
auto_mode = True
show_debug = True
manual_control = False

# 用户界面状态
slider_dragging = None  # 当前拖拽的滑块
slider_positions = {
    'p': BASE_P,
    'i': BASE_I,
    'd': BASE_D,
    'predict': PREDICTION_TIME
}


# ==================== 核心函数 ====================
def get_adaptive_factors(error):
    """根据误差大小返回自适应参数比例系数"""
    abs_error = abs(error)

    if abs_error > ADAPTIVE_PARAMS["large_error_threshold"]:
        return ADAPTIVE_PARAMS["large_error_factors"]
    elif abs_error > ADAPTIVE_PARAMS["medium_error_threshold"]:
        return ADAPTIVE_PARAMS["medium_error_factors"]
    else:
        return ADAPTIVE_PARAMS["small_error_factors"]


def advanced_pid_control(target_pos, current_pos, dt, I_past, prev_error,
                         base_p, base_i, base_d):
    """
    带抗饱和机制的高级PID控制器
    返回: (输出, 新的积分项, 当前误差, 各分量值, 使用的参数, 抗饱和状态)
    """
    if dt < 1e-6:
        dt = 1e-6

    # 计算误差
    error = target_pos - current_pos

    # 获取自适应比例系数
    p_factor, i_factor, d_factor = get_adaptive_factors(error)

    # 计算实际PID参数
    paraP = base_p * p_factor
    paraI = base_i * i_factor
    paraD = base_d * d_factor

    # 比例项
    P = error * paraP

    # 微分项（使用速度反馈，更稳定）
    if dt > 1e-3:
        D = -paraD * ball_velocity[0] * 0.7  # 使用小球自身速度，减少噪声
    else:
        D = 0

    # 积分项
    integral_input = error * paraI * dt

    # 抗饱和机制：如果输出接近限幅，减少积分累积
    estimated_output = P + I_past + D
    anti_windup_active = abs(estimated_output) > OUTPUT_LIMIT * ANTI_WINDUP_THRESHOLD

    if anti_windup_active:
        # 减少积分累积速率
        integral_input *= 0.3

    # 更新积分项
    I = I_past + integral_input
    I = max(I_MIN, min(I_MAX, I))

    # 计算总输出
    output = P + I + D

    # 输出限幅
    output = max(-OUTPUT_LIMIT, min(OUTPUT_LIMIT, output))

    return (output, I, error, P, I, D, (paraP, paraI, paraD), anti_windup_active)


def predict_position_with_acceleration(mouse_pos):
    """使用加速度模型进行轨迹预测"""
    if len(mouse_history) < 3:
        return mouse_pos[0], mouse_pos[1]

    # 计算速度和加速度
    positions = np.array(mouse_history[-3:])
    velocities = positions[1:] - positions[:-1]

    if len(velocities) >= 2:
        acceleration = velocities[1] - velocities[0]
    else:
        acceleration = np.array([0, 0])

    # 预测位置 = 当前位置 + 速度*时间 + 0.5*加速度*时间²
    current_velocity = velocities[-1] if len(velocities) > 0 else np.array([0, 0])
    predict_time = slider_positions['predict']

    prediction = (positions[-1] +
                  current_velocity * predict_time +
                  0.5 * acceleration * predict_time ** 2)

    return prediction[0], prediction[1]


def draw_slider(x, y, width, value, min_val, max_val, label):
    """绘制滑块控件"""
    # 滑块轨道
    pygame.draw.rect(screen, (200, 200, 200), (x, y, width, 8), border_radius=4)

    # 滑块位置
    ratio = (value - min_val) / (max_val - min_val) if max_val > min_val else 0
    slider_x = x + int(ratio * width)

    # 判断是否被拖拽
    mouse_pos = pygame.mouse.get_pos()
    mouse_pressed = pygame.mouse.get_pressed()[0]

    global slider_dragging
    if (slider_dragging == label and mouse_pressed) or \
            (not slider_dragging and mouse_pressed and
             abs(mouse_pos[0] - slider_x) < 15 and abs(mouse_pos[1] - y - 4) < 15):
        slider_dragging = label

    # 绘制滑块
    color = (0, 120, 255) if slider_dragging == label else (100, 100, 100)
    pygame.draw.circle(screen, color, (slider_x, y + 4), 8)

    # 标签
    text = font.render(f"{label.upper()}: {value:.3f}", True, (50, 50, 50))
    screen.blit(text, (x, y - 20))

    return slider_x


def draw_info_panel(current_fps, anti_windup_active, abs_error):
    """绘制信息面板"""
    panel_x, panel_y = 10, 10
    panel_width, panel_height = 300, 200

    # 面板背景
    pygame.draw.rect(screen, (255, 255, 255), (panel_x, panel_y, panel_width, panel_height),
                     border_radius=8)
    pygame.draw.rect(screen, (180, 180, 180), (panel_x, panel_y, panel_width, panel_height),
                     width=2, border_radius=8)

    # 标题
    title = large_font.render("PID 控制器状态", True, (0, 0, 0))
    screen.blit(title, (panel_x + 10, panel_y + 10))

    # 状态信息
    info_texts = [
        f"模式: {'自动' if auto_mode else '手动'}",
        f"实时FPS: {current_fps:.1f}",
        f"预测时间: {slider_positions['predict']:.2f}s",
        f"抗饱和: {'激活' if anti_windup_active else '未激活'}",
        f"误差: {abs_error:.1f}px"
    ]

    for i, text in enumerate(info_texts):
        color = (0, 0, 0)
        if "抗饱和" in text and anti_windup_active:
            color = COLOR_WARNING
        elif "抗饱和" in text:
            color = COLOR_SUCCESS

        info_surf = font.render(text, True, color)
        screen.blit(info_surf, (panel_x + 15, panel_y + 50 + i * 25))


# ==================== 主循环 ====================
clock = pygame.time.Clock()
current_fps = 0
anti_windup_active = False
abs_error = 0

while True:
    dt = clock.get_time() / 1000.0
    if dt < 1e-6:
        dt = 1e-6

    frame_count += 1
    total_frames += 1

    # FPS计算
    current_time = time.time()
    if current_time - last_fps_update >= 1.0:
        current_fps = frame_count / (current_time - last_fps_update)
        frame_count = 0
        last_fps_update = current_time
        fps_history.append(current_fps)
        if len(fps_history) > 60:
            fps_history.pop(0)

    # 获取鼠标位置
    mouse_pos = pygame.mouse.get_pos()

    # 更新鼠标历史
    mouse_history.append(np.array(mouse_pos))
    if len(mouse_history) > MOUSE_HISTORY_SIZE:
        mouse_history.pop(0)

    # 计算鼠标速度和加速度
    if prev_mouse_pos is not None:
        mouse_velocity = [(mouse_pos[0] - prev_mouse_pos[0]) / dt,
                          (mouse_pos[1] - prev_mouse_pos[1]) / dt]

        if len(mouse_history) >= 3:
            pos1, pos2, pos3 = mouse_history[-3:]
            vel1 = pos2 - pos1
            vel2 = pos3 - pos2
            mouse_acceleration = (vel2 - vel1) / dt

    prev_mouse_pos = mouse_pos

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                auto_mode = not auto_mode
            elif event.key == pygame.K_d:
                show_debug = not show_debug
            elif event.key == pygame.K_m:
                manual_control = not manual_control

        elif event.type == pygame.MOUSEBUTTONUP:
            slider_dragging = None

    # 清屏
    screen.fill(COLOR_BG)

    if auto_mode and not manual_control:
        # 轨迹预测
        predict_x, predict_y = predict_position_with_acceleration(mouse_pos)

        # PID控制X轴
        output_x, I_past_x, error_x, P_x, I_x, D_x, params_x, anti_windup_x = advanced_pid_control(
            predict_x, ball_pos[0], dt, I_past_x, previous_error_x,
            slider_positions['p'], slider_positions['i'], slider_positions['d']
        )

        # PID控制Y轴
        output_y, I_past_y, error_y, P_y, I_y, D_y, params_y, anti_windup_y = advanced_pid_control(
            predict_y, ball_pos[1], dt, I_past_y, previous_error_y,
            slider_positions['p'], slider_positions['i'], slider_positions['d']
        )

        # 更新抗饱和状态
        anti_windup_active = anti_windup_x or anti_windup_y
        abs_error = max(abs(error_x), abs(error_y))

        # 更新小球速度（用于D项）
        ball_velocity[0] = SMOOTH_FACTOR * ball_velocity[0] + (1 - SMOOTH_FACTOR) * output_x
        ball_velocity[1] = SMOOTH_FACTOR * ball_velocity[1] + (1 - SMOOTH_FACTOR) * output_y

        # 更新小球位置
        ball_pos[0] += ball_velocity[0]
        ball_pos[1] += ball_velocity[1]

        # 更新上一次误差
        previous_error_x = error_x
        previous_error_y = error_y

    else:
        # 手动模式：直接跟随鼠标
        ball_pos[0], ball_pos[1] = mouse_pos
        predict_x, predict_y = mouse_pos
        anti_windup_active = False
        abs_error = 0

    # 边界限制
    ball_pos[0] = max(ball_radius, min(screen_size[0] - ball_radius, ball_pos[0]))
    ball_pos[1] = max(ball_radius, min(screen_size[1] - ball_radius, ball_pos[1]))

    # ==================== 绘制 ====================

    # 绘制连接线
    pygame.draw.line(screen, COLOR_LINE_CURRENT, mouse_pos, ball_pos, 2)
    pygame.draw.line(screen, COLOR_LINE_PREDICT, mouse_pos, (int(predict_x), int(predict_y)), 2)

    # 绘制目标标记
    pygame.draw.circle(screen, COLOR_TARGET, mouse_pos, target_radius, 2)
    pygame.draw.circle(screen, COLOR_PREDICTION, (int(predict_x), int(predict_y)),
                       prediction_radius)

    # 绘制小球
    pygame.draw.circle(screen, COLOR_BALL, (int(ball_pos[0]), int(ball_pos[1])), ball_radius)

    # 绘制高亮效果
    pygame.draw.circle(screen, (255, 255, 255), (int(ball_pos[0]), int(ball_pos[1])),
                       ball_radius - 2, 1)

    # 绘制信息面板
    if show_debug:
        draw_info_panel(current_fps, anti_windup_active, abs_error)

        # 绘制滑块控件
        slider_y = 230
        slider_width = 200

        # P参数滑块
        p_val = draw_slider(50, slider_y, slider_width, slider_positions['p'],
                            0.1, 1.5, 'p')
        # I参数滑块
        i_val = draw_slider(50, slider_y + 40, slider_width, slider_positions['i'],
                            0.01, 0.3, 'i')
        # D参数滑块
        d_val = draw_slider(50, slider_y + 80, slider_width, slider_positions['d'],
                            0.05, 0.6, 'd')
        # 预测时间滑块
        predict_val = draw_slider(50, slider_y + 120, slider_width,
                                  slider_positions['predict'], 0.05, 0.3, 'predict')

        # 如果滑块被拖拽，更新值
        if slider_dragging:
            mouse_x = pygame.mouse.get_pos()[0]
            ratio = max(0, min(1, (mouse_x - 50) / slider_width))

            if slider_dragging == 'p':
                slider_positions['p'] = 0.1 + ratio * 1.4
            elif slider_dragging == 'i':
                slider_positions['i'] = 0.01 + ratio * 0.29
            elif slider_dragging == 'd':
                slider_positions['d'] = 0.05 + ratio * 0.55
            elif slider_dragging == 'predict':
                slider_positions['predict'] = 0.05 + ratio * 0.25

    # 绘制控制提示
    hints = [
        "空格: 切换自动/手动",
        "D: 显示/隐藏调试",
        "M: 切换手动控制",
        "拖拽滑块调节参数"
    ]

    for i, hint in enumerate(hints):
        hint_surf = font.render(hint, True, (100, 100, 100))
        screen.blit(hint_surf, (screen_size[0] - 180, 20 + i * 25))

    # 更新显示
    pygame.display.flip()
    clock.tick(100)